from flask import Flask, request, jsonify
from flask_cors import CORS
from flask_socketio import SocketIO, emit, join_room, leave_room
import pymysql
import random
import smtplib
from email.mime.text import MIMEText
from email.header import Header
import hashlib
import time
import jwt
import datetime
import joblib
import pandas as pd
import numpy as np
from sklearn.preprocessing import MinMaxScaler

app = Flask(__name__)
CORS(app, resources={r"/*": {"origins": "*"}})

# JWT 和 SocketIO 配置
SECRET_KEY = "your-secret-key"  # 请替换为安全的密钥
app.config['SECRET_KEY'] = SECRET_KEY
socketio = SocketIO(app, cors_allowed_origins="*")

# 固定专家的用户名
EXPERT_USERNAME = "123456"


# 数据库连接配置
def get_db_connection():
    return pymysql.connect(
        host='localhost',
        user='root',
        password='123456',
        database='dachuang',
        charset='utf8mb4',
        cursorclass=pymysql.cursors.DictCursor
    )


# 生成6位验证码
def generate_verification_code():
    return ''.join([str(random.randint(0, 9)) for _ in range(6)])


# 发送QQ邮箱验证码
def send_verification_email(email, code):
    sender = '952304392@qq.com'
    password = 'lwlbpkrftgzxbbfb'

    msg = MIMEText(f'您的验证码是：{code}，有效期20分钟', 'plain', 'utf-8')
    msg['From'] = Header(sender)
    msg['To'] = Header(email)
    msg['Subject'] = Header('验证码')

    try:
        smtp_obj = smtplib.SMTP_SSL('smtp.qq.com', 465)
        smtp_obj.login(sender, password)
        smtp_obj.sendmail(sender, [email], msg.as_string())
        smtp_obj.quit()
        return True
    except Exception as e:
        print(f"邮件发送失败: {e}")
        return False


# 密码哈希
def hash_password(password):
    return hashlib.sha256(password.encode()).hexdigest()


# 生成 JWT Token
def generate_token(user_id, username):
    payload = {
        'user_id': user_id,
        'username': username,
        'exp': datetime.datetime.utcnow() + datetime.timedelta(hours=24)
    }
    return jwt.encode(payload, SECRET_KEY, algorithm='HS256')


# 验证 JWT Token
def verify_token(token):
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=['HS256'])
        return payload
    except jwt.ExpiredSignatureError:
        return None
    except jwt.InvalidTokenError:
        return None


# 请求注册验证码
@app.route('/register/code', methods=['POST'])
def register_code():
    data = request.get_json()
    username = data.get('username')
    email = data.get('email')

    if not all([username, email]):
        return jsonify({'message': '缺少必要参数'}), 400

    conn = get_db_connection()
    c = conn.cursor()

    c.execute('SELECT * FROM users WHERE username = %s OR email = %s', (username, email))
    if c.fetchone():
        conn.close()
        return jsonify({'message': '用户名或邮箱已存在'}), 400

    code = generate_verification_code()

    if send_verification_email(email, code):
        c.execute('INSERT INTO verification_codes (email, code, timestamp) VALUES (%s, %s, %s) '
                  'ON DUPLICATE KEY UPDATE code = %s, timestamp = %s',
                  (email, code, int(time.time()), code, int(time.time())))
        conn.commit()
        conn.close()
        return jsonify({'message': '验证码已发送至您的邮箱'})
    else:
        conn.close()
        return jsonify({'message': '验证码发送失败'}), 500


# 注册接口
@app.route('/register', methods=['POST'])
def register():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')
    email = data.get('email')
    code = data.get('code')

    if not all([username, password, email, code]):
        return jsonify({'message': '缺少必要参数'}), 400

    conn = get_db_connection()
    c = conn.cursor()

    c.execute('SELECT code, timestamp FROM verification_codes WHERE email = %s', (email,))
    result = c.fetchone()

    if not result:
        conn.close()
        return jsonify({'message': '验证码无效'}), 400

    stored_code, timestamp = result['code'], result['timestamp']

    if int(time.time()) - timestamp > 1200:
        conn.close()
        return jsonify({'message': '验证码已过期'}), 400

    if stored_code != code:
        conn.close()
        return jsonify({'message': '验证码错误'}), 400

    hashed_password = hash_password(password)
    c.execute('INSERT INTO users (username, password, email, verified) VALUES (%s, %s, %s, %s)',
              (username, hashed_password, email, 1))
    conn.commit()
    c.execute('DELETE FROM verification_codes WHERE email = %s', (email,))
    conn.commit()
    conn.close()

    return jsonify({'message': '注册成功'})


# 请求登录验证码
@app.route('/login/code', methods=['POST'])
def login_code():
    data = request.get_json()
    email = data.get('email')

    if not email:
        return jsonify({'message': '缺少邮箱参数'}), 400

    conn = get_db_connection()
    c = conn.cursor()

    c.execute('SELECT * FROM users WHERE email = %s', (email,))
    if not c.fetchone():
        conn.close()
        return jsonify({'message': '邮箱未注册'}), 400

    code = generate_verification_code()

    if send_verification_email(email, code):
        c.execute('INSERT INTO verification_codes (email, code, timestamp) VALUES (%s, %s, %s) '
                  'ON DUPLICATE KEY UPDATE code = %s, timestamp = %s',
                  (email, code, int(time.time()), code, int(time.time())))
        conn.commit()
        conn.close()
        return jsonify({'message': '验证码已发送至您的邮箱'})
    else:
        conn.close()
        return jsonify({'message': '验证码发送失败'}), 500


# 登录接口
@app.route('/login', methods=['POST'])
def login():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')
    email = data.get('email')
    code = data.get('code')

    if not all([username, password, email, code]):
        return jsonify({'message': '缺少必要参数'}), 400

    conn = get_db_connection()
    c = conn.cursor()

    c.execute('SELECT id, password, email FROM users WHERE username = %s', (username,))
    result = c.fetchone()

    if not result:
        conn.close()
        return jsonify({'message': '用户不存在'}), 400

    user_id, stored_password, stored_email = result['id'], result['password'], result['email']

    if stored_email != email:
        conn.close()
        return jsonify({'message': '邮箱不匹配'}), 400

    if hash_password(password) != stored_password:
        conn.close()
        return jsonify({'message': '密码错误'}), 400

    c.execute('SELECT code, timestamp FROM verification_codes WHERE email = %s', (email,))
    result = c.fetchone()

    if not result:
        conn.close()
        return jsonify({'message': '验证码无效'}), 400

    stored_code, timestamp = result['code'], result['timestamp']

    if int(time.time()) - timestamp > 1200:
        conn.close()
        return jsonify({'message': '验证码已过期'}), 400

    if stored_code != code:
        conn.close()
        return jsonify({'message': '验证码错误'}), 400

    c.execute('DELETE FROM verification_codes WHERE email = %s', (email,))
    conn.commit()
    conn.close()

    token = generate_token(user_id, username)
    return jsonify({
        'message': '登录成功',
        'token': token,
        'user_id': user_id,
        'username': username
    })


# 获取所有用户（专家端）
@app.route('/users', methods=['GET'])
def get_users():
    token = request.headers.get('Authorization')
    if not token:
        return jsonify({'message': '缺少 token'}), 401
    if token.startswith('Bearer '):
        token = token[7:]
    payload = verify_token(token)
    if not payload:
        return jsonify({'message': '无效或过期的 token'}), 401

    if payload['username'] != EXPERT_USERNAME:
        return jsonify({'message': '仅专家可访问'}), 403

    conn = get_db_connection()
    c = conn.cursor()
    c.execute('SELECT id, username FROM users WHERE username != %s', (EXPERT_USERNAME,))
    users = c.fetchall()
    conn.close()

    return jsonify([{'id': u['id'], 'username': u['username']} for u in users])


# 获取聊天历史
@app.route('/messages/<int:receiver_id>', methods=['GET'])
def get_messages(receiver_id):
    token = request.headers.get('Authorization')
    if not token:
        return jsonify({'message': '缺少 token'}), 401
    if token.startswith('Bearer '):
        token = token[7:]
    payload = verify_token(token)
    if not payload:
        return jsonify({'message': '无效或过期的 token'}), 401

    user_id = payload['user_id']
    username = payload['username']

    conn = get_db_connection()
    c = conn.cursor()

    if username != EXPERT_USERNAME:
        c.execute('SELECT id FROM users WHERE username = %s', (EXPERT_USERNAME,))
        expert = c.fetchone()
        if not expert or receiver_id != expert['id']:
            conn.close()
            return jsonify({'message': '用户只能与专家聊天'}), 403

    c.execute('''
        SELECT sender_id, receiver_id, content, timestamp
        FROM messages
        WHERE (sender_id = %s AND receiver_id = %s) OR (sender_id = %s AND receiver_id = %s)
        ORDER BY timestamp ASC
    ''', (user_id, receiver_id, receiver_id, user_id))
    messages = c.fetchall()
    conn.close()

    return jsonify([{
        'sender_id': m['sender_id'],
        'receiver_id': m['receiver_id'],
        'content': m['content'],
        'timestamp': m['timestamp'].isoformat()
    } for m in messages])


# 风险预测接口
@app.route('/predict_risk', methods=['POST'])
def predict_risk():
    data = request.get_json()
    required_fields = [
        'debtToAssetRatio',
        'netProfitMargin',
        'cashFlowCoverageRatio',
        'receivablesToRevenueRatio',
        'operatingCashFlowRevenueRatio'
    ]

    # 验证输入
    if not all(field in data for field in required_fields):
        return jsonify({'message': '缺少必要参数'}), 400

    input_data = {k: data[k] for k in required_fields}

    try:
        # 加载模型
        model_dict = joblib.load('风控模型_v1.pkl')

        # 英文字段转中文字段（用于与模型列名匹配）
        field_mapping = {
            'debtToAssetRatio': '资产负债率',
            'netProfitMargin': '净利润率',
            'cashFlowCoverageRatio': '现金流量比率',
            'receivablesToRevenueRatio': '应收账款周转率(次)',
            'operatingCashFlowRevenueRatio': '经营净现金流/营业总收入'
        }

        # 构建 DataFrame，并重命名列为中文
        df = pd.DataFrame([input_data]).rename(columns=field_mapping)

        # 转换百分比字段
        for col in model_dict['percent_cols']:
            df[col] = df[col] / 100

        # 模型预测
        y_proba = model_dict['model'].predict_proba(df)[:, 1][0]
        y_pred = 1 if y_proba > 0.5 else 0

        # 风险指数
        scaler = MinMaxScaler(feature_range=(0, 100))
        risk_index = scaler.fit_transform([[y_proba]])[0][0]

        # 预警信息
        warning = "红色预警" if risk_index < 60 else "黄色预警" if risk_index < 70 else "正常"

        # 返回结果
        return jsonify({
            '风险概率': round(y_proba, 4),
            '风险标签': '风险' if y_pred == 1 else '安全',
            '风险指数': round(risk_index, 2),
            '预警信息': warning
        })

    except Exception as e:
        return jsonify({'message': f'预测失败：{str(e)}'}), 500


# WebSocket 事件
@socketio.on('connect')
def handle_connect():
    token = request.args.get('token')
    if not token:
        return False
    payload = verify_token(token)
    if not payload:
        return False
    emit('status', {'message': '连接成功'})


@socketio.on('join')
def on_join(data):
    token = data.get('token')
    if not token:
        emit('error', {'message': '缺少 token'})
        return
    payload = verify_token(token)
    if not payload:
        emit('error', {'message': '无效或过期的 token'})
        return

    user_id = payload['user_id']
    room = f'user_{user_id}'
    join_room(room)
    emit('status', {'message': f'用户 {user_id} 已加入房间 {room}'}, room=room)


@socketio.on('leave')
def on_leave(data):
    token = data.get('token')
    if not token:
        emit('error', {'message': '缺少 token'})
        return
    payload = verify_token(token)
    if not payload:
        emit('error', {'message': '无效或过期的 token'})
        return

    user_id = payload['user_id']
    room = f'user_{user_id}'
    leave_room(room)
    emit('status', {'message': f'用户 {user_id} 已离开房间 {room}'}, room=room)


@socketio.on('message')
def handle_message(data):
    token = data.get('token')
    if not token:
        emit('error', {'message': '缺少 token'})
        return
    payload = verify_token(token)
    if not payload:
        emit('error', {'message': '无效或过期的 token'})
        return

    sender_id = payload['user_id']
    username = payload['username']
    receiver_id = data.get('receiver_id')
    content = data.get('content')

    if not all([receiver_id, content]):
        emit('error', {'message': '缺少必要参数'})
        return

    conn = get_db_connection()
    c = conn.cursor()

    c.execute('SELECT id FROM users WHERE id = %s', (receiver_id,))
    if not c.fetchone():
        conn.close()
        emit('error', {'message': '接收者不存在'})
        return

    if username != EXPERT_USERNAME:
        c.execute('SELECT id FROM users WHERE username = %s', (EXPERT_USERNAME,))
        expert = c.fetchone()
        if not expert or receiver_id != expert['id']:
            conn.close()
            emit('error', {'message': '用户只能与专家聊天'})
            return

    c.execute('INSERT INTO messages (sender_id, receiver_id, content) VALUES (%s, %s, %s)',
              (sender_id, receiver_id, content))
    conn.commit()

    c.execute('SELECT timestamp FROM messages WHERE id = LAST_INSERT_ID()')
    timestamp = c.fetchone()['timestamp']
    conn.close()

    receiver_room = f'user_{receiver_id}'
    sender_room = f'user_{sender_id}'
    message_data = {
        'sender_id': sender_id,
        'receiver_id': receiver_id,
        'content': content,
        'timestamp': timestamp.isoformat()
    }
    emit('message', message_data, room=receiver_room)
    emit('message', message_data, room=sender_room)


if __name__ == '__main__':
    socketio.run(app, debug=True, host='0.0.0.0', port=5000, allow_unsafe_werkzeug=True)